Adding support for image special tokens in tokenizers and Message dataclass across the library. Starting with Llama3Tokenizer as an example, as multimodal Flamingo will be using this.

This PR will enabled support for multiple images per sample, interleaved throughout the text. This is done by overhauling how content is represented in the Message class:

[
    Message(role="user", content=[{"type": "image"}, {"type": "text", "content":"hello"}, ...]),
    ...
]

For backwards compatibility, passing a string into the content field will automatically get casted into the list of dictionaries. Using dictionaries is visually clunky, but it enables the most flexibility for any non-text content that needs to be tokenized differently and future metadata fields that may affect tokenization.

We must also use a different paradigm for chat formats. Now, treat it as extra tags that are prepended/appended to the Message content list based on the role:

for message in sample:
      content = (
          [{"type": "text", "content": cls.template[message.role][0]}]
          + message.content
          + [{"type": "text", "content": cls.template[message.role][1]}]
      )
      formatted_dialogue.append(
          Message(
              role=message.role,
              content=content,
              masked=message.masked,
              ipython=message.ipython,
              eot=message.eot,
          ),
      )

Changelog

• Overhaul to Message class. Update Message content field to be multimodal - can have text or image content in a list of dictionaries. Added properties for easily checking if there's image content or getting the text only content
• Update all chat format tests
• Update all chat formats to use the new Message content structure. Now we simply prepend and append tags as needed
• If Message content contains type==image, then add image special token in Llama3Tokenizer

Test Plan

Compared to a reference implementation token by token. Made this into new unit tests for text only messages, text + image, text + interleaved image, tool calls

Other considered approaches

There are other approaches I considered for enabling interleaved images in a message. The important workflow to consider is that users will have to use our utilities or create their own transforms to get their datasets into our Message format with interleaved image content. We should make this workflow as easy as possible.

1

Use an enum, Media.IMAGE, and make Message.content a List[Union[str, Media]] that combines texts and places images in between strings

[
    Message(role="user", content=[Media.IMAGE, "hello", Media.IMAGE, "world"]),
    ...
]

Pros: slightly less tedious to convert user datasets into this structure and also process this (versus dictionaries)
Cons: not easy to support image metadata or other fields that may affect image tokenization in the future

2

Keep everything as a string and allow special tokens in the text itself. Tokenizer splits on the indicated image token and adds the correct special token id and encodes the strings normally

[
    Message(role="user", content="<img> hello <img> world", media={"image": "<img>"})
    ...
]

Pros: More prompt template / chat format friendly, as we can still format the string as is. In other approaches, since the text is broken up, prompt templating is considerably more difficult. Easier to process string only content, no impact on text-only workflows
Cons: we are breaking our rule of no special tokens in the text